#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Recommendation generation functions for item-based collaborative filtering system.

Created on: 2025-04-20
Author: Nianqing Liu
"""

import numpy as np
import pandas as pd


def recommend_items(user_id, ratings_matrix, predicted_ratings, movies_df, top_n=10):
    """
    Generate top-N recommendations for a user using item-based collaborative filtering.

    Parameters:
    -----------
    user_id : int
        ID of the user to generate recommendations for (1-based).
    ratings_matrix : numpy.ndarray
        Original user-item rating matrix.
    predicted_ratings : numpy.ndarray
        Predicted rating matrix.
    movies_df : pandas.DataFrame
        DataFrame containing movie information.
    top_n : int
        Number of recommendations to generate.

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing recommended movies with predicted ratings.
    """
    # Adjust user_id to 0-based indexing
    user_idx = user_id - 1

    # Get ratings for the user
    user_ratings = ratings_matrix[user_idx]
    user_predictions = predicted_ratings[user_idx]

    # Find items the user hasn't rated yet
    unrated_items = np.where(user_ratings == 0)[0]

    # If the user has rated all items, return empty DataFrame
    if len(unrated_items) == 0:
        return pd.DataFrame(columns=["movie_id", "movie_title", "predicted_rating"])

    # Get predicted ratings for unrated items
    unrated_predictions = user_predictions[unrated_items]

    # Sort in descending order of predicted rating
    sorted_indices = np.argsort(-unrated_predictions)

    # Get top_n recommendations
    top_item_indices = sorted_indices[:top_n]
    top_items = unrated_items[top_item_indices]
    top_ratings = unrated_predictions[top_item_indices]

    # Convert to 1-based movie IDs
    movie_ids = top_items + 1

    # Create recommendations DataFrame
    recommendations = pd.DataFrame(
        {"movie_id": movie_ids, "predicted_rating": top_ratings}
    )

    # Join with movie information
    recommendations = recommendations.merge(
        movies_df[["movie_id", "movie_title"]], on="movie_id"
    )

    # Sort by predicted rating (descending)
    recommendations = recommendations.sort_values("predicted_rating", ascending=False)

    return recommendations


def recommend_for_all_users(ratings_matrix, predicted_ratings, movies_df, top_n=10):
    """
    Generate recommendations for all users using item-based collaborative filtering.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        Original user-item rating matrix.
    predicted_ratings : numpy.ndarray
        Predicted rating matrix.
    movies_df : pandas.DataFrame
        DataFrame containing movie information.
    top_n : int
        Number of recommendations per user.

    Returns:
    --------
    dict
        Dictionary mapping user IDs to recommendation DataFrames.
    """
    n_users = ratings_matrix.shape[0]
    all_recommendations = {}

    for user_idx in range(n_users):
        # Convert to 1-based user ID
        user_id = user_idx + 1

        # Generate recommendations for this user
        user_recs = recommend_items(
            user_id, ratings_matrix, predicted_ratings, movies_df, top_n
        )

        all_recommendations[user_id] = user_recs

    return all_recommendations


def get_popular_items(ratings_matrix, movies_df, top_n=10):
    """
    Get the most popular items based on number of ratings.
    Used for cold-start recommendations or fallback.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    movies_df : pandas.DataFrame
        DataFrame containing movie information.
    top_n : int
        Number of popular items to return.

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing popular movies.
    """
    # Count number of ratings for each item
    item_rating_counts = np.sum(ratings_matrix > 0, axis=0)

    # Get average rating for each item (when rated)
    rated_mask = ratings_matrix > 0
    item_rating_sums = np.sum(ratings_matrix, axis=0)
    item_rating_avgs = np.zeros_like(item_rating_counts, dtype=float)
    nonzero_counts = item_rating_counts > 0
    item_rating_avgs[nonzero_counts] = (
        item_rating_sums[nonzero_counts] / item_rating_counts[nonzero_counts]
    )

    # Create DataFrame with item stats
    item_stats = pd.DataFrame(
        {
            "movie_id": np.arange(1, len(item_rating_counts) + 1),
            "rating_count": item_rating_counts,
            "avg_rating": item_rating_avgs,
        }
    )

    # Join with movie information
    item_stats = item_stats.merge(movies_df[["movie_id", "movie_title"]], on="movie_id")

    # Sort by rating count (descending)
    item_stats = item_stats.sort_values("rating_count", ascending=False)

    # Return top_n popular items
    return item_stats.head(top_n)


def recommend_similar_items(item_id, similarity_matrix, movies_df, top_n=10):
    """
    Recommend items similar to a given item using item-based similarity.

    Parameters:
    -----------
    item_id : int
        ID of the reference item (1-based).
    similarity_matrix : numpy.ndarray
        Item-item similarity matrix.
    movies_df : pandas.DataFrame
        DataFrame containing movie information.
    top_n : int
        Number of similar items to recommend.

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing similar items.
    """
    # Adjust item_id to 0-based indexing
    item_idx = item_id - 1

    # Get similarities for the given item
    item_similarities = similarity_matrix[item_idx]

    # Set self-similarity to -1 to exclude it
    item_similarities[item_idx] = -1

    # Get indices of top similar items
    top_similar_indices = np.argsort(-item_similarities)[:top_n]
    top_similarities = item_similarities[top_similar_indices]

    # Convert to 1-based movie IDs
    movie_ids = top_similar_indices + 1

    # Create similar items DataFrame
    similar_items = pd.DataFrame(
        {"movie_id": movie_ids, "similarity": top_similarities}
    )

    # Join with movie information
    similar_items = similar_items.merge(
        movies_df[["movie_id", "movie_title"]], on="movie_id"
    )

    # Sort by similarity (descending)
    similar_items = similar_items.sort_values("similarity", ascending=False)

    return similar_items


def get_top_rated_items_for_user(
    user_id, ratings_matrix, movies_df, min_rating=4.0, top_n=5
):
    """
    Get top-rated items for a specific user.
    Useful for generating 'more like this' recommendations.

    Parameters:
    -----------
    user_id : int
        ID of the user (1-based).
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    movies_df : pandas.DataFrame
        DataFrame containing movie information.
    min_rating : float
        Minimum rating to consider an item as liked by the user (default: 4.0).
    top_n : int
        Maximum number of top-rated items to return.

    Returns:
    --------
    pandas.DataFrame
        DataFrame containing user's top-rated items.
    """
    # Adjust user_id to 0-based indexing
    user_idx = user_id - 1

    # Get user's ratings
    user_ratings = ratings_matrix[user_idx]

    # Find rated items with ratings >= min_rating
    liked_items = np.where(user_ratings >= min_rating)[0]

    if len(liked_items) == 0:
        return pd.DataFrame(columns=["movie_id", "movie_title", "user_rating"])

    # Get ratings for these items
    liked_ratings = user_ratings[liked_items]

    # Sort by rating (descending)
    sorted_indices = np.argsort(-liked_ratings)

    # Get top-n items
    top_indices = sorted_indices[: min(top_n, len(sorted_indices))]
    top_items = liked_items[top_indices]
    top_ratings = liked_ratings[top_indices]

    # Convert to 1-based movie IDs
    movie_ids = top_items + 1

    # Create top rated items DataFrame
    top_rated = pd.DataFrame({"movie_id": movie_ids, "user_rating": top_ratings})

    # Join with movie information
    top_rated = top_rated.merge(movies_df[["movie_id", "movie_title"]], on="movie_id")

    # Sort by rating (descending)
    top_rated = top_rated.sort_values("user_rating", ascending=False)

    return top_rated
